from mpi4py import MPI

def mpi_run():
    if MPI.COMM_WORLD.Get_rank() == 0:
        mpi_controller()
    else:
        mpi_worker()

def mpi_controller(f):
    '''Controls distribution of data and writes to file'''

    process = 1
    N=0

    while process < MPI.COMM_WORLD.Get_size() and N<N_tot:
        MPI.COMM_WORLD.send(N,dest=process, tag=3)
        print 'Sending Linac', N, 'to process', process
        N+=1
        process+=1

    received_processes=0
    while received_processes<N_tot:
        Results = MPI.COMM_WORLD.recv(source=MPI.ANY_SOURCE,tag=1)
        f.Write(Results['Linac'])
        process = Results['rank']
        received_processes+=1

        if N<N_tot:
            MPI.COMM_WORLD.send(dest=process,tag=3)
            print 'Sending Linac', N, 'to process', process
            N+=1

    for process in range(1,MPI.COMM_WORLD.size):
        MPI.COMM_WORLD.send(-1,dest=process,tag=3)



def mpi_worker():
    '''Worker process'''

    while True:
        n = MPI.COMM_WORLD.recv(source=0,tag=3)
        if n=-1: break

        do_something()

        results = {'rank':rank,'Linac':Lin}
        MPI.COMM_WORLD.send(results,dest=0,tag=1)

